# file:D:\share\python\python_net\base_model\nlp\transformer1\UsherConfig.py
import os
from typing import Literal

from transformers import PretrainedConfig, AutoTokenizer


class UsherConfig(PretrainedConfig):
    model_type = "usherTransformer"
    auto_map = {
        "AutoModelForCausalLM": "Transformer.Transformer"
    }

    def __init__(
            self,
            path=r"result",
            # 基本参数
            vocab_size=1024,
            dim=1024,
            n_layers=8,
            normal_alpha_init=0.5,
            inter_dim=2048,
            # MLA
            n_heads=16,
            qk_head_dim=2048,

            # MOE
            n_dense_layers=1,
            n_routed_experts=8,
            n_shared_experts=1,
            n_active_experts=1,
            moe_inter_dim=2048,
            **kwargs
    ):
        """

     Attributes:
         vocab_size (`int`, *optional*, defaults to 1024): 分词长度
         dim (`int`, *optional*, defaults to 1024): 词向量维度
         n_layers (`int`, *optional*, defaults to 8): 层数
         dinamic_tanh 初始alpha参数 (`float`, *optional*, defaults to 0.5): 正态分布初始化参数
         path (`str`, *optional*, defaults to r"result"): 模型保存路径
         normal_alpha_init (`float`, *optional*, defaults to 0.5): 正态分布初始化参数
         inter_dim (`int`, *optional*, defaults to 2048): 前馈神经网络中间维度

         n_heads (`int`, *optional*, defaults to 16): 多头注意力头数
         qk_head_dim (`int`, *optional*, defaults to 2048): qk维度

         n_dense_layers (`int`, *optional*, defaults to 1): 正常前馈神经网络层数(超出则是MOE)
         n_routed_experts (`int`, *optional*, defaults to 8): 每层专家数量
         n_shared_experts (`int`, *optional*, defaults to 1): 共享专家数量
         n_active_experts (`int`, *optional*, defaults to 1): 激活专家数量
         moe_inter_dim (`int`, *optional*, defaults to 2048): MOE中间维度


    """
        # 调用父类初始化方法
        super().__init__(**kwargs)
        # 初始化模型参数
        self.vocab_size = vocab_size
        self.dim = dim
        self.n_layers = n_layers
        self.normal_alpha_init = normal_alpha_init
        self.n_heads = n_heads
        self.n_dense_layers = n_dense_layers
        self.qk_head_dim = qk_head_dim
        self.inter_dim = inter_dim
        self.n_routed_experts = n_routed_experts
        self.moe_inter_dim = moe_inter_dim
        self.n_shared_experts = n_shared_experts
        self.n_active_experts = n_active_experts
        # 转换为绝对路径
        self.path = os.path.abspath(path)
